{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-LLM Agentic Sales Agent\n",
    "\n",
    "This notebook details the following:\n",
    "1. Different models: Demonstrates how to integrate and use various language models (DeepSeek, Gemini, Llama3.3,  Claude, Amazon Bedrock Nova) within the agentic framework.\n",
    "2. Structured Outputs: Shows how to define and enforce structured outputs for agents using Pydantic BaseModel.\n",
    "3. Guardrails: Implements an input guardrail to prevent sensitive information (like personal names) from being processed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing necessary components from Strands SDK and other libraries.\n",
    "from dotenv import load_dotenv\n",
    "from openai import AsyncOpenAI\n",
    "from agents import Agent, Runner, trace, function_tool, OpenAIChatCompletionsModel, input_guardrail, GuardrailFunctionOutput\n",
    "from typing import Dict\n",
    "import sendgrid\n",
    "import os\n",
    "from sendgrid.helpers.mail import Mail, Email, To, Content\n",
    "from pydantic import BaseModel\n",
    "from agents.extensions.models.litellm_model import LitellmModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load environment variables from a .env file. `override=True` will overwrite existing environment variables.\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve API keys from environment variables.\n",
    "amazon_bedrock_api_key = os.getenv('AMAZON_BEDROCK_API_KEY')\n",
    "google_api_key = os.getenv('GEMINI_API_KEY')\n",
    "deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')\n",
    "groq_api_key = os.getenv('GROQ_API_KEY')\n",
    "\n",
    "# Print the status of each API key to confirm they are loaded.\n",
    "if amazon_bedrock_api_key:\n",
    "    print(f\"Amazon Bedrock API Key exists and begins {amazon_bedrock_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"Amazon Bedrock API Key not set\")\n",
    "\n",
    "if google_api_key:\n",
    "    print(f\"Google API Key exists and begins {google_api_key[:2]}\")\n",
    "else:\n",
    "    print(\"Google API Key not set (and this is optional)\")\n",
    "\n",
    "if deepseek_api_key:\n",
    "    print(f\"DeepSeek API Key exists and begins {deepseek_api_key[:3]}\")\n",
    "else:\n",
    "    print(\"DeepSeek API Key not set (and this is optional)\")\n",
    "\n",
    "if groq_api_key:\n",
    "    print(f\"Groq API Key exists and begins {groq_api_key[:4]}\")\n",
    "else:\n",
    "    print(\"Groq API Key not set (and this is optional)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define different instruction sets for various sales agent personas.\n",
    "# These instructions guide the behavior and tone of the generated emails.\n",
    "\n",
    "# Professional sales agent instructions\n",
    "professional_instructions = \"You are a professional sales agent working for ComplAI, \\\n",
    "a company that provides a SaaS tool for ensuring SOC2 compliance and preparing for audits, powered by AI. \\\n",
    "You write professional, serious cold emails.\"\n",
    "\n",
    "# Engaging sales agent instructions\n",
    "engaging_instructions = \"You are a humorous, engaging sales agent working for ComplAI, \\\n",
    "a company that provides a SaaS tool for ensuring SOC2 compliance and preparing for audits, powered by AI. \\\n",
    "You write witty, engaging cold emails that are likely to get a response.\"\n",
    "\n",
    "# Concise sales agent instructions\n",
    "concise_instructions = \"You are a busy sales agent working for ComplAI, \\\n",
    "a company that provides a SaaS tool for ensuring SOC2 compliance and preparing for audits, powered by AI. \\\n",
    "You write concise, to the point cold emails.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- It's easy to use any models with OpenAI compatible endpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define base URLs for different OpenAI-compatible API endpoints.\n",
    "GEMINI_BASE_URL = \"https://generativelanguage.googleapis.com/v1beta/openai/\"\n",
    "DEEPSEEK_BASE_URL = \"https://api.deepseek.com/v1\"\n",
    "GROQ_BASE_URL = \"https://api.groq.com/openai/v1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize AsyncOpenAI clients for each service.\n",
    "# These clients will be used to interact with the respective LLM APIs.\n",
    "deepseek_client = AsyncOpenAI(base_url=DEEPSEEK_BASE_URL, api_key=deepseek_api_key)\n",
    "gemini_client = AsyncOpenAI(base_url=GEMINI_BASE_URL, api_key=google_api_key)\n",
    "groq_client = AsyncOpenAI(base_url=GROQ_BASE_URL, api_key=groq_api_key)\n",
    "\n",
    "# Create OpenAIChatCompletionsModel instances for each model.\n",
    "# These models abstract the API interactions and provide a consistent interface for the Agents SDK.\n",
    "deepseek_model = OpenAIChatCompletionsModel(model=\"deepseek-chat\", openai_client=deepseek_client)\n",
    "gemini_model = OpenAIChatCompletionsModel(model=\"gemini-2.0-flash\", openai_client=gemini_client)\n",
    "llama3_3_model = OpenAIChatCompletionsModel(model=\"llama-3.3-70b-versatile\", openai_client=groq_client)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize Agent instances for each sales persona, assigning them a specific model and instructions.\n",
    "professional_sales_agent = Agent(name=\"DeepSeek Sales Agent\", \n",
    "                                instructions=professional_instructions, \n",
    "                                model=deepseek_model)\n",
    "engaging_sales_agent =  Agent(name=\"Gemini Sales Agent\", \n",
    "                                instructions=engaging_instructions, \n",
    "                                model=gemini_model)\n",
    "concise_sales_agent  = Agent(name=\"Llama3.3 Sales Agent\",\n",
    "                                instructions=concise_instructions, \n",
    "                                model=llama3_3_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a common description for the sales email writing tools.\n",
    "description = \"Write a cold sales email\"\n",
    "\n",
    "# Convert each sales agent into a tool.\n",
    "# This allows a higher-level agent (e.g., Sales Manager) to invoke these agents as functions.\n",
    "professional_tool = professional_sales_agent.as_tool(tool_name=\"professional_sales_agent\", \n",
    "                                                    tool_description=description)\n",
    "engaging_tool = engaging_sales_agent.as_tool(tool_name=\"engaging_sales_agent\", \n",
    "                                            tool_description=description)\n",
    "concise_tool = concise_sales_agent.as_tool(tool_name=\"concise_sales_agent\", \n",
    "                                            tool_description=description)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function tool to send HTML emails using SendGrid.\n",
    "@function_tool\n",
    "def send_html_email(subject: str, html_body: str) -> Dict[str, str]:\n",
    "    \"\"\" Send out an email with the given subject and HTML body to all sales prospects \"\"\"\n",
    "    print(\"[Tool Call] Send html tool engaged\")\n",
    "    sg = sendgrid.SendGridAPIClient(api_key=os.environ.get('SENDGRID_API_KEY'))\n",
    "    from_email = Email(\"iankisali@gmail.com\")\n",
    "    to_email = To(\"iankisali295@gmail.com\")\n",
    "    content = Content(\"text/html\", html_body)\n",
    "    mail = Mail(from_email, to_email, subject, content).get()\n",
    "    response = sg.client.mail.send.post(request_body=mail)\n",
    "    print(\"[Tool Call] Send html tool completed\")\n",
    "    return {\"status\": \"success\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define instructions for agents responsible for writing email subjects and converting text to HTML.\n",
    "\n",
    "# subject writer instructions\n",
    "subject_instructions = \"You can write a subject for a cold sales email. \\\n",
    "You are given a message and you need to write a subject for an email that is likely to get a response.\"\n",
    "\n",
    "# html body writer instructions\n",
    "html_instructions = \"You can convert a text email body to an HTML email body. \\\n",
    "You are given a text email body which might have some markdown \\\n",
    "and you need to convert it to an HTML email body with simple, clear, compelling layout and design.\"\n",
    "\n",
    "# Initialize an Agent for writing email subjects, using a LiteLLMModel for Anthropic Claude.\n",
    "subject_writer = Agent(name=\"Email subject writer\", \n",
    "                        instructions=subject_instructions, \n",
    "                        model=LitellmModel(model=\"anthropic.claude-3-5-sonnet-20240620-v1:0\", api_key=amazon_bedrock_api_key))\n",
    "\n",
    "# Convert the subject writer agent into a tool.\n",
    "subject_tool = subject_writer.as_tool(tool_name=\"subject_writer_tool\", \n",
    "                        tool_description=\"Write a subject for a cold sales email\")\n",
    "\n",
    "# Initialize an Agent for converting email bodies to HTML, also using Anthropic Claude.\n",
    "html_converter = Agent(name=\"HTML email body converter\", \n",
    "                        instructions=html_instructions, \n",
    "                        model=LitellmModel(model=\"anthropic.claude-3-5-sonnet-20240620-v1:0\", api_key=amazon_bedrock_api_key))\n",
    "\n",
    "# Convert the HTML converter agent into a tool.\n",
    "html_tool = html_converter.as_tool(tool_name=\"html_converter_tool\",\n",
    "                        tool_description=\"Convert a text email body to an HTML email body\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Group all email-related tools together.\n",
    "email_tools = [subject_tool, html_tool, send_html_email]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the list of email tools (for verification).\n",
    "email_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define instructions for the Email Manager agent.\n",
    "email_formater_instructions =\"You are an email formatter and sender. You receive the body of an email to be sent. \\\n",
    "You first use the subject_writer tool to write a subject for the email, then use the html_converter tool to convert the body to HTML. \\\n",
    "Finally, you use the send_html_email tool to send the email with the subject and HTML body.\"\n",
    "\n",
    "# Initialize the Email Manager agent. This agent orchestrates the email formatting and sending process.\n",
    "emailer_manager_agent = Agent(\n",
    "    name=\"Email Manager\",\n",
    "    instructions=email_formater_instructions,\n",
    "    tools=email_tools,\n",
    "    model=LitellmModel(model=\"amazon.nova-pro-v1:0\", api_key=amazon_bedrock_api_key),\n",
    "    handoff_description=\"Convert an email to HTML and send it\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Group all sales agent tools.\n",
    "tools = [professional_tool, concise_tool, engaging_tool]\n",
    "# Define handoffs for the Sales Manager, which is the Email Manager agent.\n",
    "handoffs = [emailer_manager_agent]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define instructions for the Sales Manager agent.\n",
    "sales_manager_instructions = \"You are a sales manager working for ComplAI. You use the tools given to you to generate cold sales emails. \\\n",
    "You never generate sales emails yourself; you always use the tools. \\\n",
    "You try all 3 sales agent tools at least once before choosing the best one. \\\n",
    "You can use the tools multiple times if you're not satisfied with the results from the first try. \\\n",
    "You select the single best email using your own judgement of which email will be most effective. \\\n",
    "After picking the email, you handoff to the Email Manager agent to format and send the email.\"\n",
    "\n",
    "# Initialize the Sales Manager agent. This is the top-level agent that decides which sales agent to use and then hands off to the Email Manager.\n",
    "sales_manager = Agent(\n",
    "    name=\"Sales Manager\",\n",
    "    instructions=sales_manager_instructions,\n",
    "    tools=tools,\n",
    "    handoffs=handoffs,\n",
    "    model=LitellmModel(model=\"amazon.nova-pro-v1:0\", api_key=amazon_bedrock_api_key))\n",
    "\n",
    "# Define the initial message for the sales manager.\n",
    "message = \"Send out a cold sales email addressed to Dear Safaricom CEO from Ian Kisali\"\n",
    "\n",
    "# Run the sales process and trace its execution.\n",
    "# The `trace` context manager helps in visualizing the agent's decision-making process.\n",
    "with trace(\"Automated SDR\"):\n",
    "    result = await Runner.run(sales_manager, message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Check out the trace on OpenAI SDK Developer platform:\n",
    "\n",
    "https://platform.openai.com/traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a Pydantic BaseModel for structured output from the name check guardrail.\n",
    "class NameCheckOutput(BaseModel):\n",
    "    is_name_in_message: bool\n",
    "    name: str\n",
    "\n",
    "# Initialize a guardrail agent. Its purpose is to check if a personal name is present in the input message.\n",
    "guardrail_agent = Agent( \n",
    "    name=\"Name checker agent\",\n",
    "    instructions=\"Check if the user is including someone's personal name in what they want you to do.\",\n",
    "    output_type=NameCheckOutput,\n",
    "    model=LitellmModel(model=\"anthropic.claude-3-5-sonnet-20240620-v1:0\", api_key=amazon_bedrock_api_key)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define an input guardrail function.\n",
    "# This function uses the `guardrail_agent` to check for names and triggers a \"tripwire\" if a name is found.\n",
    "# If a name is found, `tripwire_triggered` is set to True, which will halt the execution.\n",
    "@input_guardrail\n",
    "async def name_guardrail(ctx, agent, message):\n",
    "    result = await Runner.run(guardrail_agent, message, context=ctx.context)\n",
    "    is_name_in_message = result.final_output.is_name_in_message\n",
    "    return GuardrailFunctionOutput(output_info={\"found_name\": result.final_output},tripwire_triggered=is_name_in_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a \"careful\" sales manager agent that includes the `name_guardrail`.\n",
    "# This agent will prevent execution if a personal name is detected in the initial message.\n",
    "careful_sales_manager = Agent(\n",
    "    name=\"Sales Manager\",\n",
    "    instructions=sales_manager_instructions,\n",
    "    tools=tools,\n",
    "    handoffs=[emailer_manager_agent],\n",
    "    model=deepseek_model,\n",
    "    input_guardrails=[name_guardrail]\n",
    "    )\n",
    "\n",
    "# Attempt to run with a message containing a personal name.\n",
    "# This is expected to trigger the guardrail and raise an error\n",
    "message = \"Send out a cold sales email addressed to Dear Safaricom CEO from Ian Kisali\"\n",
    "\n",
    "with trace(\"Protected Automated SDR\"):\n",
    "    result = await Runner.run(careful_sales_manager, message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Check out the trace:\n",
    "\n",
    "https://platform.openai.com/traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Attempt to run with a message that does not contain a personal name.\n",
    "# This should proceed without triggering the guardrail.\n",
    "message = \"Send out a cold sales email addressed to Dear CEO from Head of Business Development\"\n",
    "\n",
    "with trace(\"Protected Automated SDR\"):\n",
    "    result = await Runner.run(careful_sales_manager, message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
